import * as tfvis from '@tensorflow/tfjs-vis'
import * as tf from '@tensorflow/tfjs'

/**
 * 归一化
 */
window.onload = () => {
	const heights = [150, 160, 170]
	const weights = [40, 50, 60]

	// 散点图
	tfvis.render.scatterplot(
	{
		name: "身高体重训练集"
	}, {
		values: heights.map((x, i) => ({x, y: weights[i]}))
	}, {
		xAxisDomain: [140, 180],
		yAxisDomain: [30, 70]
	})

	const inputs = tf.tensor(heights).sub(150).div(20)
	const labels = tf.tensor(weights).sub(40).div(20)

	/**
	 * 连续的模型
	 */
	const model = tf.sequential()

	/** 全链接层
	 * units 神经元数
	 * inputShape 输入形状
	 */
	model.add(tf.layers.dense({
		units: 1, inputShape: [1]
	}))

	/**
	 * 设置损失函数、优化器sgd(学习率)
	 */
	model.compile({
		loss: tf.losses.meanSquaredError,
		optimizer: tf.train.sgd(0.1)
	})

	/**fit 开始学习任务
	 * batchSize 取样数
	 * epochs 对训练数据数组进行迭代的次数。
	 */
	model.fit(inputs, labels, {
		batchSize: 3,
		epochs: 200,
		callbacks: tfvis.show.fitCallbacks({
			name: "训练过程"
		}, [
			'loss'
		])
	})

	// debugger
	window.predict = (form) => {
		const output = model.predict(tf.tensor([form.x.value * 1]).sub(150).div(20));
		alert(`如果 身高 为 ${form.x.value}，那么预测 体重 为 ${output.mul(20).add(40).dataSync()[0]}kg`);
	}
}

